/*
 * The MIT License (MIT)
 *
 * Copyright (c) 2019 Code Technology Studio
 *
 * Permission is hereby granted, free of charge, to any person obtaining a copy of
 * this software and associated documentation files (the "Software"), to deal in
 * the Software without restriction, including without limitation the rights to
 * use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of
 * the Software, and to permit persons to whom the Software is furnished to do so,
 * subject to the following conditions:
 *
 * The above copyright notice and this permission notice shall be included in all
 * copies or substantial portions of the Software.
 *
 * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS
 * FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
 * COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER
 * IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
 * CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
 */
package org.dromara.jpom.service.manage;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.convert.Convert;
import cn.hutool.core.util.StrUtil;
import cn.hutool.core.util.URLUtil;
import cn.hutool.extra.ssh.ChannelType;
import cn.hutool.extra.ssh.JschUtil;
import com.alibaba.fastjson2.JSONObject;
import com.alibaba.fastjson2.JSONValidator;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.jcraft.jsch.*;
import org.dromara.jpom.service.manage.webshell.ConstantPool;
import org.dromara.jpom.service.manage.webshell.SSHConnectInfo;
import org.dromara.jpom.service.manage.webshell.WebSSHData;
import org.dromara.jpom.util.SocketSessionUtil;
import org.dromara.jpom.util.StringUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;

/**
 * @Description: WebSSH业务逻辑实现
 * @Author: NoCortY
 * @Date: 2020/3/8
 */
@Service
public class WebSSHService {
    @Autowired
    private WebSSHConfig webSSHConfig;
    //存放ssh连接信息的map
    private static Map<String, Object> sshMap = new ConcurrentHashMap<>();

    private Logger logger = LoggerFactory.getLogger(WebSSHService.class);
    //线程池
    private ExecutorService executorService = Executors.newCachedThreadPool();


    /**
     * @Description: 初始化连接
     * @Param: [session]
     * @return: void
     * @Author: NoCortY
     * @Date: 2020/3/7
     */
    public void initConnection(javax.websocket.Session session) {
        JSch jSch = new JSch();
        SSHConnectInfo sshConnectInfo = new SSHConnectInfo();
        sshConnectInfo.setjSch(jSch);
        sshConnectInfo.setWebSocketSession(session);
        String userId = this.getParameters(session, ConstantPool.USER_UUID_KEY);
        //将这个ssh连接信息放入map中
        sshMap.put(userId, sshConnectInfo);
        logger.info("ssh connect init, userId = {}", userId);
    }

    private String getParameters(javax.websocket.Session session, String name) {
        Map<String, List<String>> requestParameterMap = session.getRequestParameterMap();
        Map<String, String> parameters = session.getPathParameters();
        List<String> strings = requestParameterMap.get(name);
        String value = CollUtil.join(strings, StrUtil.COMMA);
        if (StrUtil.isEmpty(value)) {
            value = parameters.get(name);
        }
        return URLUtil.decode(value);
    }

    /**
     * @Description: 处理客户端发送的数据
     * @Param: [buffer, session]
     * @return: void
     * @Author: NoCortY
     * @Date: 2020/3/7
     */
    public void recvHandle(String buffer, javax.websocket.Session session) {
        ObjectMapper objectMapper = new ObjectMapper();
        WebSSHData webSSHData = null;
        try {
            webSSHData = objectMapper.readValue(buffer, WebSSHData.class);
        } catch (IOException e) {
            logger.error("Json转换异常");
            logger.error("异常信息:{}", e.getMessage());
            return;
        }
        String userId = this.getParameters(session, ConstantPool.USER_UUID_KEY);
        if (ConstantPool.WEBSSH_OPERATE_CONNECT.equals(webSSHData.getOperate())) {
            //找到刚才存储的ssh连接对象
            SSHConnectInfo sshConnectInfo = (SSHConnectInfo) sshMap.get(userId);
            //启动线程异步处理
            WebSSHData finalWebSSHData = webSSHData;

            finalWebSSHData.setHost(webSSHConfig.getHost());
            finalWebSSHData.setPort(webSSHConfig.getPort());
            finalWebSSHData.setUsername(webSSHConfig.getUsername());
            finalWebSSHData.setPassword(webSSHConfig.getPassword());

            executorService.execute(new Runnable() {
                public void run() {
                    try {
                        connectToSSH(sshConnectInfo, finalWebSSHData, session);
                    } catch (JSchException | IOException e) {
                        logger.error("webssh连接异常");
                        logger.error("异常信息:{}", e.getMessage());
                        try {
                            sendMessage(session, "在agent上开启ssh连接失败，server到agent之间的ws关闭".getBytes(StandardCharsets.UTF_8));
                        } catch (IOException ex) {
                            logger.error("发送消息异常:{}", e.getMessage());
                        }
                        close(session);
                    }
                }
            });
        } else if (ConstantPool.WEBSSH_OPERATE_COMMAND.equals(webSSHData.getOperate())) {
            String command = webSSHData.getCommand();
            SSHConnectInfo sshConnectInfo = (SSHConnectInfo) sshMap.get(userId);

            if (sshConnectInfo != null) {
                try {
                    JSONValidator.Type type = StringUtil.validatorJson(command);
                    if (type == JSONValidator.Type.Object) {
                        JSONObject resizeInfo = JSONObject.parseObject(command);
                        String data = resizeInfo.getString("data");
                        if (StrUtil.equals(data, "jpom-heart")) {
                            // 心跳消息不转发
                            return;
                        }

                        if (StrUtil.equals(data, "resize")) {
                            // 缓存区大小
                            if (sshConnectInfo.getChannel() == null || !sshConnectInfo.getChannel().isConnected()) {
                                executorService.submit(() -> {
                                    int i = 0;
                                    while (i++ < 10) {
                                        try {
                                            Thread.sleep(1000);
                                            if(sshConnectInfo.getChannel() != null){
                                                this.resize((ChannelShell) sshConnectInfo.getChannel(), resizeInfo);
                                                logger.info("ssh 连接建立,开始执行 resize");
                                                break;
                                            }else{
                                                logger.info("ssh 连接还没建立,继续等待建立后进行resize, 已经等待 {} 次", i);
                                            }
                                        } catch (Exception e) {
                                            logger.error("定时监听ssh连接建立异常",e);
                                        }
                                    }
                                });
                            }else{
                                this.resize((ChannelShell) sshConnectInfo.getChannel(), resizeInfo);
                            }
                        }
                    } else {
                        transToSSH(sshConnectInfo.getChannel(), command);
                    }
                } catch (IOException e) {
                    logger.error("webssh连接异常");
                    logger.error("异常信息:{}", e.getMessage());
                    close(session);
                }
            }
        } else {
            logger.error("不支持的操作");
            close(session);
        }
    }


    private void resize(ChannelShell channel, JSONObject jsonObject) {
        Integer rows = Convert.toInt(jsonObject.getString("rows"), 10);
        Integer cols = Convert.toInt(jsonObject.getString("cols"), 10);
        Integer wp = Convert.toInt(jsonObject.getString("wp"), 10);
        Integer hp = Convert.toInt(jsonObject.getString("hp"), 10);
        channel.setPtySize(cols, rows, wp, hp);
        logger.info("resize执行完成");
    }

    public void sendMessage(javax.websocket.Session session, byte[] buffer) throws IOException {
        String msg = new String(buffer);
        if (logger.isDebugEnabled()) {
            logger.debug(msg);
        }
        SocketSessionUtil.send(session, msg);
    }


    public void close(javax.websocket.Session session) {
        String userId = this.getParameters(session, ConstantPool.USER_UUID_KEY);
        SSHConnectInfo sshConnectInfo = (SSHConnectInfo) sshMap.get(userId);
        if (sshConnectInfo != null) {
            //断开连接
            if (sshConnectInfo.getChannel() != null) {
                sshConnectInfo.getChannel().disconnect();
            }

            //map中移除
            sshMap.remove(userId);
            logger.info("ssh connect close, userId = {}", userId);
        }
    }


    /**
     * @Description: 使用jsch连接终端
     * @Param: [cloudSSH, webSSHData, webSocketSession]
     * @return: void
     * @Author: NoCortY
     * @Date: 2020/3/7
     */
    private void connectToSSH(SSHConnectInfo sshConnectInfo, WebSSHData webSSHData, javax.websocket.Session webSocketSession) throws JSchException, IOException {
        Session sshSession = null;
        Properties config = new Properties();
        config.put("StrictHostKeyChecking", "no");
        //获取jsch的会话
        sshSession = sshConnectInfo.getjSch().getSession(webSSHData.getUsername(), webSSHData.getHost(), webSSHData.getPort());
        sshSession.setConfig(config);
        //设置密码
        sshSession.setPassword(webSSHData.getPassword());
        //连接  超时时间30s
        sshSession.connect(30000);

        //开启shell通道
        Channel sshChannel = sshSession.openChannel("shell");

        //通道连接 超时时间3s
        sshChannel.connect(3000);

        //设置channel
        sshConnectInfo.setChannel(sshChannel);

        //转发消息
        transToSSH(sshChannel, "\r");

        //读取终端返回的信息流
        InputStream inputStream = sshChannel.getInputStream();
        try {
            //循环读取
            byte[] buffer = new byte[1024];
            int i = 0;
            //如果没有数据来，线程会一直阻塞在这个地方等待数据。
            while ((i = inputStream.read(buffer)) != -1) {
                sendMessage(webSocketSession, Arrays.copyOfRange(buffer, 0, i));
            }

        } finally {
            //断开连接后关闭会话
            sshSession.disconnect();
            sshChannel.disconnect();
            if (inputStream != null) {
                inputStream.close();
            }
        }

    }

    /**
     * @Description: 将消息转发到终端
     * @Param: [channel, data]
     * @return: void
     * @Author: NoCortY
     * @Date: 2020/3/7
     */
    private void transToSSH(Channel channel, String command) throws IOException {
        if (channel != null) {
            OutputStream outputStream = channel.getOutputStream();
            outputStream.write(command.getBytes());
            outputStream.flush();
        }
    }

}
